import pandas as pd
import numpy as np
from typing import Tuple, overload
from numpy import array
from pandas import DataFrame, date_range
from torch.utils.data import Dataset


class Slicer(Dataset):
    """
    Class for slicing data before feature selector and model constructor.
    Parameters:
        parameters used in both slicers before feature selector and model constructor:
            x, y
            parameters about index:
                freq: 'D', 'W', 'M', 'Q'
                end_datetime, periods, test_len
            parameters about columns:
                x columns: features, features = None when instantiating this class before feature selector
                y columns: y_name
        parameters used only in slicer before model constructor (that's why we only give these three parameters defaults):
            x_time_step, y_time_step, valid_pct
            prediction: False, high frequency;
                        True, prediction
    """

    def __init__(self, x: DataFrame, y: DataFrame, y_name: list, freq: str, end_datetime: str, periods: int, test_len: int,
                 features: list, n_features: int = 5, x_time_step: int = 1, prediction: bool = False, y_time_step: int = 1,
                 valid_pct: float = 0.2) -> None:
        self.x_time_step = x_time_step
        self.y_time_step = y_time_step

        x_idx = date_range(end=end_datetime, periods=periods, freq=freq)
        if prediction:
            y_end_datetime = date_range(start=end_datetime, periods=y_time_step + 1, freq=freq)[-1]
            darnn_y_idx = date_range(end=y_end_datetime, periods=periods+y_time_step, freq=freq)
        else:
            y_end_datetime = date_range(start=end_datetime, periods=y_time_step, freq=freq)[-1]
            darnn_y_idx = date_range(end=y_end_datetime, periods=periods, freq=freq)
        y_idx = date_range(end=y_end_datetime, periods=periods + y_time_step - x_time_step, freq=freq)
        

        self.x = x.loc[x_idx[: -self.x_time_step * test_len], features[: n_features]] if features else x.loc[x_idx[: -self.x_time_step * test_len], :]
        self.y = y.loc[y_idx[: -self.y_time_step * test_len], y_name]

        self._x = x.loc[x_idx, features[: n_features]] if features else x.loc[x_idx, :]
        self._y = y.loc[y_idx, y_name]
        
        self.xx = x.loc[x_idx, features[: n_features]] if features else x.loc[x_idx, :]
        self.yy = y.loc[darnn_y_idx, y_name]

        splitting_point = round((self.__len__() - test_len) * (1 - valid_pct))
        self.train_set = self.__getitem__(slice(0, splitting_point, 1))
        self.valid_set = self.__getitem__(slice(splitting_point, self.__len__() - test_len, 1))
        self.test_set = self.__getitem__(slice(self.__len__() - test_len, self.__len__(), 1))

        self.x_train = [item[0] for item in self.train_set]
        self.y_train = [item[1] for item in self.train_set]
        self.x_valid = [item[0] for item in self.valid_set]
        self.y_valid = [item[1] for item in self.valid_set]
        self.x_test = [item[0] for item in self.test_set]
        self.y_test = [item[1] for item in self.test_set]

        self.dataset = (self.x_train, self.y_train, self.x_valid, self.y_valid, self.x_test)

    @overload
    def __getitem__(self, idx: int) -> Tuple[array, array]:
        pass

    @overload
    def __getitem__(self, idx: slice) -> list:
        pass

    def __getitem__(self, idx):
        if isinstance(idx, int):
            if 0 <= idx < self.__len__():
                return self._x.values[idx: idx + self.x_time_step, :], \
                       self._y.values[idx: idx + self.y_time_step]
            else:
                """or return self.x[idx:, :]"""
                raise IndexError('Index is out of range.')
        elif isinstance(idx, slice):
            start, stop, step = idx.indices(self.__len__())
            return [self.__getitem__(i) for i in range(start, stop, step)]

    def __len__(self) -> int:
        """Notice that it is the length of dataset not the length of x or y!"""
        return self._x.shape[0] - self.x_time_step + 1


class Slicer_for_time(Dataset):

    def __init__(self, freq: str, end_datetime: str, periods: int, x_time_step: int = 1, y_time_step: int = 1, test_len : int = 1,
                 valid_pct: float = 0.2, prediction: bool = False) -> None:
        self.freq = freq
        self.periods = periods
        self.x_time_step = x_time_step
        self.y_time_step = y_time_step
        self.test_len = test_len
        self.prediction = prediction
        if self.prediction:
            self.end_datetime = date_range(start = end_datetime, periods = self.y_time_step + self.test_len, freq=self.freq).strftime('%F')[-1]
        else:
            self.end_datetime = date_range(start = end_datetime, periods = self.y_time_step + self.test_len - 1, freq=self.freq).strftime('%F')[-1]
        x_idx = date_range(end=self.end_datetime, periods=self.periods, freq=self.freq)
        x_idx = [str(t)[:10] for t in list(x_idx)]
        self.times = pd.DataFrame(np.nan, index=x_idx[(self.x_time_step - 1):], columns=['value'])
        self.splitting_point = round((self.__len__() - self.y_time_step) * (1 - valid_pct))

        train_times = self.__getitem__(slice(0, self.splitting_point, 1))
        self.train_times = self.listof_times(train_times)
        valid_times = self.__getitem__(slice(self.splitting_point, self.__len__() - self.y_time_step, 1))
        self.valid_times = self.listof_times(valid_times)
        test_times = self.__getitem__(slice(self.__len__() - self.y_time_step - 1, self.__len__(), 1))
        self.test_times = self.listof_times(test_times)

    @overload
    def __getitem__(self, idx: int) -> str:
        pass

    @overload
    def __getitem__(self, idx: slice) -> list:
        pass

    def __getitem__(self, idx):
        if isinstance(idx, int):
            return self.times.index.values[idx: idx + self.x_time_step]
        elif isinstance(idx, slice):
            start, stop, step = idx.indices(self.__len__())
            return [self.__getitem__(i) for i in range(start, stop, step)]

    def listof_times(self, input_times):
        if len(input_times[0]) == 1:
            return ["".join(i.tolist()) for i in input_times]
        else:
            return [i[0] for i in input_times]

    def __len__(self) -> int:
        """Notice that it is the length of dataset not the length of x or y!"""
        return self.times.shape[0]


if __name__ == '__main__':
    xx = pd.read_pickle('../../input/high_freq_all_data_2021-03-10.pkl').iloc[:, :3]
    yy = pd.read_excel('../../input/y集合.xlsx', index_col=0).iloc[:, :]
    ff = list(xx.columns)
    # slicer = Slicer(x=xx, y=yy, freq='M', end_datetime='2018-11-30', periods=12, test_len=1, n_features=3,
    #                features=ff, y_name=['COM'], x_time_step=2, prediction=  False, y_time_step=1, valid_pct=0.2)
    slicer_for_time = Slicer_for_time(freq='M', end_datetime='2020-08-31', periods=28, x_time_step=1,
                                    y_time_step=2, valid_pct=0.2)
